from sqlalchemy.ext.compiler import compiles
import pandas as pd
import time
from sqlalchemy import create_engine, Table, MetaData
from sqlalchemy.sql.expression import Insert
from sqlalchemy.dialects.mysql import insert
from sqlalchemy.orm import sessionmaker
from sqlalchemy import create_engine, text
import platform
import json
import requests
import uuid
from urllib.parse import quote_plus
import traceback
from datetime import datetime
import re

def check_exhibition_info(exhibition_data):
    """
    调用AI接口检查展会信息，使用流式输出
    """
    
    # 生成唯一的dataId和chatId
    dataId = str(uuid.uuid4()).replace("-", "")
    chatId = str(uuid.uuid4()).replace("-", "")
    responseChatItemId = str(uuid.uuid4()).replace("-", "")
    
    
    print(exhibition_data)
    promot = f"""
    以上是报表信息，请你根据知识库查询，核对以下展会信息：
    {exhibition_data}
            仅仅返回json格式数据（其它数据不要)
            [                "status": True/fasle,
                "error_check_description": "错误信息",
                "advice_check_description": "建议信息"
            ]
            注意： 报表有些信息式正确的只是翻译成中文，这种小错误请忽略掉，
            如果知识库里未提及的而报表有的也忽略掉错误，如果没有大致错误，返回式true,error_check_description
            则是"无发现明显错误" 仅仅子啊advice_check_description提出一些建议性检查。
            如果发现明显错误（比如几乎和知识库的信息相差太大），则返回false,error_check_description则
            指出明显错误出来。
            """

    # 构建提示信息
    messages = [{
        "dataId": dataId,
        "hideInUI": False,
        "role": "user",
        "content": f"{promot}"
    }]
    
    # 使用新的API密钥和Bearer认证
    API_KEY = "fastgpt-kspcWF3r4XnI9leHXcZqE6tkY2LKGx7p9WgT3fMD9Ywtn06zyQ36oFdo4Vb"
    
    # 使用新的API接口和请求格式
    url = "http://106.53.97.117:3000/api/v1/chat/completions"
    
    payload = {
        "chatId": chatId,
        "appId": "67e014c77a46e4ccc286d102",
        "variables": {
            "cTime": time.strftime("%Y-%m-%d %H:%M:%S %A")
        },
        "messages": [
            {
                "content": promot,
                "role": "user"
            }
        ],
        "modules": [
            {
                "moduleId": "dataset",
                "datasetId": "67e3b1d27a46e4ccc2879bfd"  # 使用正确的数据集ID
            }
        ],
        "stream": True
    }
    
    headers = {
        "Accept": "text/event-stream",
        "Content-Type": "application/json",
        "Authorization": f"Bearer {API_KEY}",  # 使用Bearer令牌认证
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/134.0.0.0 Safari/537.36"
    }
    
    # 添加重试机制
    max_retries = 3
    retry_delay = 2  # 秒
    last_exception = None
    
    for retry_count in range(max_retries):
        try:
            # 增加超时设置
            response = requests.post(
                url, 
                json=payload, 
                headers=headers, 
                timeout=(10, 300),  # 连接超时10秒，读取超时300秒
                stream=True
            )
            response.raise_for_status()
            
            full_response = ""
            print(f"API响应状态码: {response.status_code}")
            
            for line in response.iter_lines():
                if line:
                    try:
                        # 移除 "data: " 前缀并解析 JSON
                        if line.startswith(b'data: '):
                            json_str = line[6:].decode('utf-8')
                            
                            # 处理结束标记
                            if json_str.strip() == '[DONE]':
                                print("收到完成标记")
                                break
                                
                            # 解析JSON数据
                            try:
                                data = json.loads(json_str)
                                
                                # 提取内容 - 新的响应格式
                                if 'data' in data and 'text' in data['data']:
                                    content = data['data']['text']
                                    print(content, end='', flush=True)
                                    full_response += content
                                # 兼容旧格式
                                elif 'choices' in data and len(data['choices']) > 0:
                                    content = data['choices'][0].get('delta', {}).get('content', '')
                                    if content:
                                        print(content, end='', flush=True)
                                        full_response += content
                            except json.JSONDecodeError:
                                # 如果无法解析为JSON，跳过
                                continue
                    except Exception as e:
                        print(f"处理流式响应时出错: {str(e)}")
                        continue
            
            print(f"\n完整响应: {full_response[:200]}...")
            
            # 验证响应是否为空
            if not full_response.strip():
                print("警告: 收集到的完整响应为空")
                if retry_count < max_retries - 1:
                    wait_time = retry_delay * (2 ** retry_count)
                    print(f"等待 {wait_time} 秒后重试...")
                    time.sleep(wait_time)
                    continue
                return {"status": False, "error": "API返回了空响应内容"}
                
            return {"status": True, "response": full_response}
            
        except requests.exceptions.RequestException as e:
            last_exception = e
            print(f"API请求失败 (尝试 {retry_count+1}/{max_retries}): {str(e)}")
            
            # 如果不是最后一次尝试，则等待后重试
            if retry_count < max_retries - 1:
                wait_time = retry_delay * (2 ** retry_count)  # 指数退避
                print(f"等待 {wait_time} 秒后重试...")
                time.sleep(wait_time)
            else:
                print(f"达到最大重试次数 ({max_retries})，放弃请求")
                return {"status": False, "error": str(e)}
    
    # 如果所有重试都失败
    return {"status": False, "error": str(last_exception) if last_exception else "所有重试都失败"}

class MySQLHelper:
    def __init__(self):
        if platform.system() == 'Linux':
            self.host = "106.53.97.117"
            self.port = 3306
            self.user = "my_user"
            self.password = "my_password"
            self.db = "my_database"
        else:
            self.host = "106.53.97.117"
            self.port = 3306
            self.user = "my_user"
            self.password = "my_password"
            self.db = "my_database"
        self.charset = "utf8mb4"
        self.engine = None
        self.Session = None
        self._create_engine()

    def _create_engine(self):
        # 创建数据库连接引擎，使用连接池
        connection_string = f"mysql+pymysql://{self.user}:{self.password}@{self.host}:{self.port}/{self.db}?charset={self.charset}"
        self.engine = create_engine(connection_string, pool_size=10, max_overflow=1490, pool_recycle=3600)
        self.Session = sessionmaker(bind=self.engine)

    def get_session(self):
        # 获取数据库会话
        return self.Session()

    def mysql_replace_into(self, table, conn, keys, data_iter):
        # 定义一个编译器扩展，用于将 INSERT INTO 替换为 REPLACE INTO
        @compiles(Insert)
        def replace_string(insert, compiler, **kw):
            s = compiler.visit_insert(insert, **kw)
            s = s.replace("INSERT INTO", "REPLACE INTO")
            return s

        # 准备数据
        data = [dict(zip(keys, row)) for row in data_iter]
        
        # 执行插入操作
        conn.execute(table.table.insert(), data)



    def _df_res_to_db(self, pd, table_name):
        pd.to_sql(
            name=table_name,
            con=self.engine.connect(),
            index=False,
            chunksize=2000,
            if_exists='append', method=self.mysql_replace_into)






    def execute_select(self, sql):
        time_start = time.time()
        with self.engine.connect() as conn:
            # 使用 text() 函数将 SQL 语句转换为可执行对象
            sql_text = text(sql)
            result = conn.execute(sql_text)
            data = []
            for row in result:
                # 使用 row._mapping 将 RowProxy 转换为字典
                data.append(dict(row._mapping))
            time_end = time.time()
            # print(f"执行时间：{time_end - time_start}")
            return data

    def execute_value(self, sql):
        with self.engine.connect() as conn:
            result = conn.execute(sql)
            return result.fetchone()

    def execute_edit(self, sql):
        with self.engine.connect() as conn:
            result = conn.execute(sql)
            conn.commit()
            return result.rowcount
    def execute_charts(self, sql):
        """
        :return 图表数据
        """
        with self.engine.connect() as conn:
            result = conn.execute(sql)
            target = {}
            for j in range(len(result.keys())):
                data = [row[j] for row in result]
                target[result.keys()[j]] = data
            return target

    def execute_datatable(self, sql):
        with self.engine.connect() as conn:
            result = conn.execute(sql)
            datas = []
            for row in result:
                target = list(row)
                datas.append(target)
            data = {"data": datas}
            return data

    def execute_cloud(self, sql):
        with self.engine.connect() as conn:
            result = conn.execute(sql)
            data = {row[0]: row[1] for row in result}
            return data


db = MySQLHelper()

fields_mapping = {
    'chinese_name': '中文名称',
    'english_name': '英文全称',
    'english_short_name': '英文简称',
    'logo': '标识',
    'holding_time': '举办时间',
    'event_date_start': '举办开始日期',
    'event_date_end': '举办结束日期',
    'event_month': '举办月份',
    'continent': '洲',
    'country': '国别',
    'city': '城市',
    'frequency': '举办频率',
    'industry': '行业',
    'organizer': '主办方',
    'organizer_en': '主办方英文',
    'contact_info': '联系方式',
    'exhibition_hall': '展馆',
    'exhibition_hall_english': '展馆英文',
    'official_website': '官网',
    'exhibit_scope': '展品范围',
    'description': '简介'
}

def get_current_status():
    """
    获取当前展会记录的处理状态统计
    """
    try:
        # 查询各状态的展会数量
        status_sql = """
            SELECT 
                SUM(CASE WHEN ai_link_status = 1 AND (ai_checked_status = 0 OR ai_checked_status IS NULL) THEN 1 ELSE 0 END) as pending_count,
                SUM(CASE WHEN ai_checked_status = 1 THEN 1 ELSE 0 END) as success_count,
                SUM(CASE WHEN ai_checked_status = 2 THEN 1 ELSE 0 END) as needs_correction_count,
                SUM(CASE WHEN ai_checked_status = 3 THEN 1 ELSE 0 END) as failed_count,
                COUNT(*) as total_count
            FROM exhibitions
        """
        
        result = db.execute_select(status_sql)
        if result and len(result) > 0:
            status_data = result[0]
            
            print("\n" + "="*50)
            print("当前展会处理状态统计")
            print("="*50)
            print(f"待处理展会记录: {status_data.get('pending_count', 0)} 条")
            print(f"AI检查通过展会: {status_data.get('success_count', 0)} 条")
            print(f"需要修正展会: {status_data.get('needs_correction_count', 0)} 条")
            print(f"检查失败展会: {status_data.get('failed_count', 0)} 条")
            print(f"总展会记录数: {status_data.get('total_count', 0)} 条")
            print("="*50)
            
            return status_data
        else:
            print("获取当前状态统计失败")
            return None
    except Exception as e:
        print(f"获取当前状态统计时出错: {str(e)}")
        print(traceback.format_exc())
        return None

def main():
    try:
        print("="*50)
        print("开始处理展会AI检查")
        print("="*50)
        
        # 获取处理前的状态统计
        print("处理前状态统计:")
        get_current_status()
        
        # 修改查询语句，只选择我们需要的字段
        fields = ', '.join(fields_mapping.keys())
        fields += ', id'  # 添加id字段用于更新
        
        sql = f"""
            SELECT {fields} 
            FROM exhibitions 
            WHERE ai_link_status = 1 
            AND (ai_checked_status = 0 OR ai_checked_status IS NULL)
            LIMIT 100  -- 添加限制，避免一次处理太多记录
        """
        
        # 尝试多次查询数据库，防止临时连接问题
        max_db_retries = 3
        result_json = None
        
        for retry in range(max_db_retries):
            try:
                result_json = db.execute_select(sql)
                break
            except Exception as e:
                print(f"数据库查询失败 (尝试 {retry+1}/{max_db_retries}): {str(e)}")
                if retry < max_db_retries - 1:
                    print("等待 5 秒后重试...")
                    time.sleep(5)
                else:
                    print("达到最大重试次数，无法查询数据库")
                    raise
        
        if not result_json:
            print("没有需要处理的展会记录")
            return
            
        result_df = pd.DataFrame(result_json)

        # 重命名列名为中文，但保留id列
        id_column = result_df['id'].copy()
        result_df = result_df.rename(columns=fields_mapping)
        result_df['id'] = id_column

        # 处理所有记录
        total_records = len(result_df)
        print(f"总共需要处理 {total_records} 条展会记录")
        
        success_count = 0
        failed_count = 0
        start_time = time.time()

        for index, row in result_df.iterrows():
            try:
                exhibition_data = row.to_dict()
                exhibition_id = exhibition_data.pop('id')  # 移除id以避免干扰AI判断
                
                print(f"\n[{index+1}/{total_records}] 正在处理第 {index+1} 条记录")
                print(f"展会ID: {exhibition_id}")
                print(f"展会名称: {exhibition_data.get('中文名称', '未知')}")
                
                # 增加处理单条记录的超时机制
                try:
                    # 调用AI接口检查展会信息
                    check_result = check_exhibition_info(exhibition_data)
                    print(f"AI处理状态: {'成功' if check_result['status'] else '失败'}")
                except Exception as e:
                    print(f"调用AI接口时出错: {str(e)}")
                    failed_count += 1
                    _update_check_failed(exhibition_id, f"AI接口调用失败: {str(e)[:100]}")
                    continue

                if check_result['status']:
                    try:
                        # 从response中提取数据
                        response_text = check_result['response']
                        
                        # 尝试从响应中提取JSON
                        json_data = None
                        
                        # 如果响应包含markdown代码块，提取其中的JSON
                        if '```json' in response_text:
                            # 提取```json和```之间的内容
                            json_start = response_text.find('```json\n') + 8
                            if json_start == 7:  # 如果没有找到```json\n，尝试不带换行符的版本
                                json_start = response_text.find('```json') + 7
                            json_end = response_text.rfind('```')
                            if json_start != -1 and json_end != -1 and json_start < json_end:
                                json_str = response_text[json_start:json_end].strip()
                                # 解析JSON数据
                                try:
                                    json_data = json.loads(json_str)
                                except:
                                    print("从Markdown代码块提取JSON失败")
                        
                        # 如果上面的方法失败，尝试在文本中搜索JSON对象
                        if not json_data:
                            # 尝试匹配JSON对象样式的内容
                            json_matches = re.findall(r'({[\s\S]*?})', response_text)
                            for match in json_matches:
                                try:
                                    potential_json = json.loads(match)
                                    # 验证是否包含我们需要的字段
                                    if 'status' in potential_json:
                                        json_data = potential_json
                                        break
                                except:
                                    continue
                        
                        # 如果仍然没有找到有效JSON，尝试直接解析整个响应
                        if not json_data:
                            try:
                                # 清理响应文本，移除可能导致JSON解析失败的字符
                                cleaned_response = response_text.strip()
                                
                                # 尝试寻找第一个{和最后一个}来提取可能的JSON对象
                                first_brace = cleaned_response.find('{')
                                last_brace = cleaned_response.rfind('}')
                                
                                if first_brace != -1 and last_brace != -1 and first_brace < last_brace:
                                    potential_json_str = cleaned_response[first_brace:last_brace+1]
                                    json_data = json.loads(potential_json_str)
                                else:
                                    # 最后尝试整个文本
                                    json_data = json.loads(cleaned_response)
                            except:
                                print("直接解析整个响应失败")
                                raise json.JSONDecodeError("无法从响应中提取有效JSON", response_text, 0)
                        
                        # 确保JSON数据包含必要的字段，或设置默认值
                        if not json_data:
                            print("无法从AI响应中解析出有效的JSON数据")
                            raise ValueError("无法从AI响应中解析出有效的JSON数据")
                        
                        # 提取状态和描述信息，并设置默认值
                        ai_status = json_data.get('status', False)
                        
                        # 处理error_check_description，将字典转换为字符串
                        error_description = json_data.get('error_check_description', '')
                        if isinstance(error_description, dict):
                            error_description = '; '.join([f"{k}: {v}" for k, v in error_description.items()])
                        error_description = str(error_description).replace("'", "''")  # 正确的SQL转义
                        
                        # 处理advice_check_description，将字典转换为字符串
                        advice_description = json_data.get('advice_check_description', '')
                        if isinstance(advice_description, dict):
                            advice_description = '; '.join([f"{k}: {v}" for k, v in advice_description.items()])
                        advice_description = str(advice_description).replace("'", "''")  # 正确的SQL转义
                        
                        # 限制字段长度，避免数据库字段长度限制问题
                        error_description = error_description[:1000] if error_description else ''
                        advice_description = advice_description[:1000] if advice_description else ''
                        
                        # 更新数据库，添加ai_checked_time
                        current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                        update_sql = f"""
                        UPDATE exhibitions 
                        SET ai_checked_status = {1 if ai_status else 2},
                            error_check_description = '{error_description}',
                            advice_check_description = '{advice_description}',
                            ai_checked_time = '{current_time}'
                        WHERE id = {exhibition_id}
                        """
                        
                        # 使用事务执行更新，并添加重试机制
                        update_success = False
                        max_update_retries = 3
                        
                        for update_retry in range(max_update_retries):
                            try:
                                with db.engine.connect() as conn:
                                    with conn.begin():
                                        conn.execute(text(update_sql))
                                update_success = True
                                break
                            except Exception as e:
                                print(f"数据库更新失败 (尝试 {update_retry+1}/{max_update_retries}): {str(e)}")
                                if update_retry < max_update_retries - 1:
                                    print("等待 3 秒后重试更新...")
                                    time.sleep(3)
                        
                        if update_success:
                            print(f"数据库更新成功 - AI状态: {'通过' if ai_status else '需要修正'}")
                            if not ai_status:
                                print(f"错误描述: {error_description}")
                                print(f"建议: {advice_description}")
                            
                            success_count += 1
                        else:
                            print("达到最大重试次数，无法更新数据库")
                            failed_count += 1
                            continue
                        
                    except json.JSONDecodeError as e:
                        print(f"JSON解析错误: {str(e)}")
                        failed_count += 1
                        # 标记检查状态为失败
                        _update_check_failed(exhibition_id, f"JSON解析错误: {str(e)[:100]}")
                        continue
                    except Exception as e:
                        print(f"处理AI检查结果时出错: {str(e)}")
                        print(traceback.format_exc())
                        failed_count += 1
                        # 标记检查状态为失败
                        _update_check_failed(exhibition_id, f"处理错误: {str(e)[:100]}")
                        continue
                else:
                    print(f"AI检查失败: {check_result.get('error', '未知错误')}")
                    failed_count += 1
                    # 标记检查状态为失败
                    _update_check_failed(exhibition_id, f"AI检查失败: {check_result.get('error', '未知错误')[:100]}")
                    continue
                
                # 每处理10条记录显示一次进度统计
                if (index + 1) % 10 == 0 or (index + 1) == total_records:
                    elapsed_time = time.time() - start_time
                    avg_time = elapsed_time / (index + 1)
                    remaining_count = total_records - (index + 1)
                    est_remaining_time = avg_time * remaining_count
                    
                    print(f"\n--- 处理进度 ---")
                    print(f"已处理: {index+1}/{total_records} ({(index+1)/total_records*100:.1f}%)")
                    print(f"成功: {success_count}, 失败: {failed_count}")
                    print(f"已用时间: {elapsed_time:.1f}秒")
                    if remaining_count > 0:
                        print(f"预计剩余时间: {est_remaining_time:.1f}秒 (约 {est_remaining_time/60:.1f}分钟)")
                
                # 添加短暂延时，避免请求过于频繁
                time.sleep(2)
                
            except Exception as e:
                print(f"处理记录时发生错误: {str(e)}")
                print(traceback.format_exc())
                failed_count += 1
                try:
                    _update_check_failed(exhibition_id, f"处理错误: {str(e)[:100]}")
                except:
                    print("更新失败状态时也发生错误")
                continue
        
        # 最终统计
        total_time = time.time() - start_time
        print("\n"+"="*50)
        print("处理完成统计")
        print("="*50)
        print(f"总记录数: {total_records}")
        print(f"成功处理: {success_count} ({success_count/total_records*100:.1f}%)")
        print(f"处理失败: {failed_count} ({failed_count/total_records*100:.1f}%)")
        print(f"总耗时: {total_time:.1f}秒 (约 {total_time/60:.1f}分钟)")
        print("="*50)
        
        # 重新查询一次当前状态，显示实时数据
        print("\n处理后状态统计:")
        get_current_status()
        
        # 检查是否还有待处理的记录
        current_status = get_current_status()
        if current_status and current_status.get('pending_count', 0) > 0:
            print(f"\n仍有 {current_status.get('pending_count')} 条记录待处理")
            print("您可以再次运行此脚本继续处理剩余记录")
        
    except Exception as e:
        print(f"程序执行过程中发生错误: {str(e)}")
        print(traceback.format_exc())
    finally:
        print("\n所有展会记录处理完成！")

def _update_check_failed(exhibition_id, reason):
    """更新展会AI状态为失败"""
    try:
        current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        safe_reason = reason.replace("'", "''")  # SQL转义
        update_sql = f"""
        UPDATE exhibitions 
        SET ai_checked_status = 3,
            error_check_description = '检查失败: {safe_reason}',
            ai_checked_time = '{current_time}'
        WHERE id = {exhibition_id}
        """
        with db.engine.connect() as conn:
            with conn.begin():
                conn.execute(text(update_sql))
        print(f"展会ID {exhibition_id} 标记为检查失败: {reason}")
        return True
    except Exception as e:
        print(f"更新展会ID {exhibition_id} 检查状态失败: {str(e)}")
        return False

if __name__ == "__main__":
    main()